from datasets import *

def test(model, enc_input, start_symbol):
    # Starting Reference: http://nlp.seas.harvard.edu/2018/04/03/attention.html#greedy-decoding
    enc_outputs, enc_self_attns = model.Encoder(enc_input)
    # 创建一个全0的张量dec_input，它的维度为[1, tgt_len]，数据类型与enc_input相同。tgt_len是目标语言句子的最大长度。
    # .type_as()方法将张量的数据类型转换为另一个张量的数据类型。
    dec_input = torch.zeros(1, tgt_len).type_as(enc_input.data)
    next_symbol = start_symbol
    # 模拟预测时 Decoder 的逐步解码过程
    for i in range(0, tgt_len):
        # 将下一个符号设置为当前位置的输出
        dec_input[0][i] = next_symbol
        dec_outputs, _, _ = model.Decoder(dec_input, enc_input, enc_outputs)
        # dec_outputs: [batch_size, tgt_len, d_model]
        # 将解码器的输出映射到目标词汇表的维度上，得到映射后的张量projected。
        projected = model.projection(dec_outputs)
        # projected: [batch_size, tgt_len, tgt_vocab_size]
        # NOTE:
        #  squeeze() 方法会删除大小为1的维度，如果指定参数，则只会删除指定维度大小为1的维度。
        #  如果要删除所有大小为1的维度，可以使用 squeeze() 或者 squeeze(dim=None)。
        #  如果不想删除任何维度，可以使用 unsqueeze() 方法增加一个维度。
        # 对projected进行处理，首先使用squeeze(0)将其维度为1的批次维度压缩掉，然后在最后一个维度上取最大值，得到最大值的索引，即为当前位置的输出符号。
        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
        # 将当前位置的输出符号设置为下一个符号。
        next_word = prob.data[i]
        next_symbol = next_word.item()
    return dec_input


enc_inputs, dec_inputs, dec_outputs = make_data()
loader = Data.DataLoader(MyDataSet(enc_inputs, dec_inputs, dec_outputs), 2, True)
enc_inputs, _, _ = next(iter(loader))
# NOTE: .pth扩展名用于表示模型参数，而.pt扩展名则用于表示包含模型和优化器状态的完整检查点。因此，.pt文件比.pth文件更大一些。
model = torch.load('model.pt')
# 预测时只给定编码器的输入以及 start_symbol
# enc_inputs：[batch_size,src_len]
predict_dec_input = test(model, enc_inputs[0].view(1, -1).cuda(), start_symbol=tgt_vocab["S"])
predict, _, _, _ = model(enc_inputs[0].view(1, -1).cuda(), predict_dec_input)
predict = predict.data.max(1, keepdim=True)[1]

print([src_idx2word[int(i)] for i in enc_inputs[0]], '->',[idx2word[n.item()] for n in predict.squeeze()])
